package com.chengv.redismq.core;

import cn.hutool.core.lang.Assert;
import com.chengv.redismq.constant.RedisMQConstant;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;

import java.time.Duration;
import java.time.Instant;
import java.util.Arrays;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledFuture;

/**
 * 延迟消息调度任务容器
 * <p>
 * 可通过 {@link #register(String)} 方法注册调度任务，每个调度任务会定时获取延迟队列中到期的消息，然后将消息发送到对应的 Stream 中
 *
 * @author chengv
 * @see RedisMQDelayMessageScheduledTaskRegistrar
 */
@Slf4j
public class RedisMQDelayMessageScheduledTaskContainer implements DisposableBean {

    private final RedisTemplate<String, String> redisTemplate;
    private TaskScheduler taskScheduler;
    private int batchSize = 100;
    private Duration pollTime = Duration.ofSeconds(2);
    private final ConcurrentHashMap<String, ScheduledFuture<?>> scheduledTasks = new ConcurrentHashMap<>();

    public RedisMQDelayMessageScheduledTaskContainer(RedisTemplate<String, String> redisTemplate) {
        this.redisTemplate = redisTemplate;
        this.taskScheduler = getDefaultTaskScheduler();
    }

    private ThreadPoolTaskScheduler getDefaultTaskScheduler() {
        ThreadPoolTaskScheduler taskScheduler = new ThreadPoolTaskScheduler();
        taskScheduler.initialize();
        return taskScheduler;
    }

    public void setTaskScheduler(TaskScheduler taskScheduler) {
        this.taskScheduler = taskScheduler;
    }

    public void setBatchSize(int batchSize) {
        Assert.isTrue(batchSize > 0, "batchSize 必须大于0");

        this.batchSize = batchSize;
    }

    public void setPollTime(Duration pollTime) {
        Assert.isTrue(!pollTime.isNegative(), "pollTime 不能是负数");

        this.pollTime = pollTime;
    }

    public void register(String topic) {
        Assert.notBlank(topic, "topic 不能为空");
        String delayQueueName = RedisMQConstant.DELAY_QUEUE_PREFIX + topic;

        // 如果任务已经注册，则忽略
        if (scheduledTasks.containsKey(delayQueueName)) return;

        // 创建一个固定速率的定时任务，每 pollTimeout 执行一次
        ScheduledFuture<?> future = taskScheduler.scheduleAtFixedRate(new ScheduledTask(delayQueueName), Instant.now(),
                pollTime);
        scheduledTasks.put(delayQueueName, future);

        log.info("[registerRedisMQDelayMessageScheduledTask][成功], topic=({})", topic);
    }

    @Override
    public void destroy() {
        scheduledTasks.values().forEach(scheduledFuture ->
                scheduledFuture.cancel(false));
    }

    private class ScheduledTask implements Runnable {
        private final String delayQueueName;
        private final String queueName;

        ScheduledTask(String delayQueueName) {
            this.delayQueueName = delayQueueName;
            this.queueName = delayQueueName.replace(RedisMQConstant.DELAY_QUEUE_PREFIX, RedisMQConstant.QUEUE_PREFIX);
        }

        @Override
        public void run() {
            // 使用 Lua 脚本，避免并发问题
            // 从延迟队列中获取到期的消息，将消息发送到对应的 Stream，然后从延迟队列中删除消息
            redisTemplate.execute(new DefaultRedisScript<>(
                            "local expiredValues = redis.call('zrangebyscore', KEYS[1], 0, ARGV[1], 'LIMIT', 0, ARGV[2]);"
                                    + "if #expiredValues > 0 then "
                                    + "    for i, v in ipairs(expiredValues) do "
                                    + "        redis.call('xadd', KEYS[2], '*', 'payload', v);"
                                    + "    end "
                                    + "    redis.call('zrem', KEYS[1], unpack(expiredValues));"
                                    + "end;",
                            Long.class),
                    Arrays.asList(delayQueueName, queueName),
                    String.valueOf(System.currentTimeMillis()), String.valueOf(batchSize));
        }
    }
}
